package com.bleeth.flow.step.output;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.poi.excel.BigExcelWriter;
import cn.hutool.poi.excel.ExcelUtil;
import com.bleeth.flow.core.common.APlugin;
import com.bleeth.flow.core.common.PluginAnnotation;
import com.bleeth.flow.core.common.PluginTypeEnum;
import com.jd.platform.async.wrapper.WorkerWrapper;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @author ：Bleeth
 * @date ：2021-08-06
 * @description：
 */
@Data
@Slf4j
@PluginAnnotation(name = "Excel写",
        type = PluginTypeEnum.OUTPUT,
        description = "",
        id = "ExcelWriterPlugin")
public class ExcelWriterPlugin extends APlugin implements Serializable {

    private static final long serialVersionUID = 9151516393920337644L;

    private String path;


    @Override
    public Dataset<Row> action(Dataset<Row> param, Map<String, WorkerWrapper> allWrappers) {
        super.action(param, allWrappers);

        String fromPluginName = fromList.get(0);
        WorkerWrapper fromWrapper = allWrappers.get(fromPluginName);
        Dataset<Row> ds = (Dataset<Row>) fromWrapper.getWorkResult().getResult();
        //通过工具类创建writer
        BigExcelWriter writer = ExcelUtil.getBigWriter(path);

        //获取字段
        StructType schema = ds.schema();
        List<String> headerList = CollUtil.newArrayList();
        for (StructField field : schema.fields()) {
            String name = field.name();
            headerList.add(name);
        }

        //转list
        List<Row> rowList = ds.collectAsList();
        System.err.println("需要写入的数据量有：" + rowList.size());


        //写表头
        List<List<String>> excelRowList = new ArrayList<>();
        excelRowList.add(headerList);
        System.err.println("写入表头完成");

        for (int i = 0; i < rowList.size(); i++) {
            Row row = rowList.get(i);
            List<String> list = new ArrayList<>();
            for (String header : headerList) {
                String value = row.getAs(header);
                list.add(value);
            }
            excelRowList.add(list);

            if (excelRowList.size() >= 1000) {
                writer.write(excelRowList);
                excelRowList = CollUtil.newArrayList();
            }
        }
        if (excelRowList.size() > 0) {
            writer.write(excelRowList);
        }
        System.err.println("数据写入excel完成...");
        writer.close();
        return null;
    }


}
